#!/usr/bin/env python
# author gesturewei, ihep, cas


def usage():
    print("Usage: check_cpu_info -w <percent>[%] -c <percent>[%]");


def main():
    EXIT_OK=0;
    EXIT_WARNING=1;
    EXIT_CRITICAL=2;
    EXIT_UNKNOWN=3;
    import sys, psutil, getopt;
    warning_value=0.0;
    critical_value=0.0;
    is_warning_set=False;
    is_critical_set=False;
    try:
        opts, args=getopt.getopt(sys.argv[1:],"w:c:h");
    except getopt.GetoptError:
        usage();
        sys.exit(EXIT_UNKNOWN);
    for key,value in opts:
        if key=='-w':
            if value.endswith('%'):
                warning_value=float(value[:-1]);
            else:
                warning_value=float(value);
            is_warning_set=True;
        elif key=='-c':
            if value.endswith('%'):
                critical_value=int(value[:-1]);
            else:
                critical_value=int(value);
            is_critical_set=True;
        elif key=='-h':
            usage();
            sys.exit(EXIT_OK);
    if not is_warning_set or not is_critical_set:
        usage();
        sys.exit(EXIT_UNKNOWN);
    if warning_value < 0.0 or critical_value < 0.0:
        print("The value must be positive.");
        sys.exit(EXIT_UNKNOWN);
    cpu_info=psutil.cpu_times_percent(.5);
    cpu_user=cpu_info[0];
    cpu_sys=cpu_info[2];
    cpu_idle=cpu_info[3];
    cpu_iowait=cpu_info[4];
    cpu_used=100-cpu_idle;
    if cpu_used>critical_value:
        out_msg="CPU CRITICAL - ";
        exit_code=EXIT_CRITICAL;
    elif cpu_used>warning_value:
        out_msg="CPU WARNING - ";
        exit_code=EXIT_WARNING;
    else:
        out_msg="CPU OK - ";
        exit_code=EXIT_OK;
    out_msg+="%.02f%% used (user: %.02f; system: %.02f; iowait: %.02f) |cpu=" % (cpu_used, cpu_user,cpu_sys,cpu_iowait);
    out_msg+="%.02f;%.02f;%.02f;0;100.00" % (cpu_used,warning_value,critical_value);
    print(out_msg);
    sys.exit(exit_code);

if __name__=="__main__":
    main();
